import torch
from utils.util import LD_param as get_res
from utils.util import AverageMeter,save_csv,mergetotxt,SSIM,PSNR,normal
from models.cyclegan.testapi import NewAverageMeter,get_frequency,draw_Fea
import os
import torch.nn as nn
import numpy as np

def change_fre(F):
        ans = 100 - F
        if ans <= 0:
            ans = 0
        return ans
 
class sinTestApi():
    def __init__(self,opt):
        # self.D_generation = opt.D_generation
        # self.frequency = opt.frequency
        self.batch_size = opt.batch_size
        self.use_inception = opt.use_inception
        self.loss = nn.MSELoss()
        self.controller = opt.controller
        self.name = opt.name
        self.version = opt.version
        self.phase = opt.phase
        self.epoch = opt.epoch
        self.loss = nn.MSELoss()
        self.outputs = opt.outputs
        self.c = 1e-9
        pass
    def init(self,mode):
        self.mode = mode

        # self.Fmode = Fmode
        if mode:
            # self.count_mean_G=AverageMeter()
            self.count_G=AverageMeter()
        else:
            self.count_mean_SB_PSNR=AverageMeter()
            self.count_mean_SAB_PSNR=AverageMeter()
            self.count_mean_SB_SSIM=NewAverageMeter()
            self.count_mean_SAB_SSIM=AverageMeter()
            
            if self.use_inception:
                self.features = {}
                self.features['real_SA'] = []
                self.features['real_SB'] = []
                self.features['fake_SB'] = []
                self.count_mean_SB_FID=NewAverageMeter()
                self.count_mean_SAB_FID=AverageMeter()
                self.count_mean_disFID=NewAverageMeter()
            
        # if self.frequency:
        self.count_Frequency_SB=NewAverageMeter()
        self.count_DFrequency_SB=NewAverageMeter()

            
    def record(self,records):
        out = {}
        if self.mode:
            self.count_mean_G.update(records['loss_G'].contiguous().detach().cpu(),self.batch_size)
            self.count_mean_G_S01.update(records['loss_G_S01'].contiguous().detach().cpu(),self.batch_size)
            self.count_mean_G_S10.update(records['loss_G_S10'].contiguous().detach().cpu(),self.batch_size)
            self.count_mean_G_S02.update(records['loss_G_S02'].contiguous().detach().cpu(),self.batch_size)
            self.count_mean_G_S20.update(records['loss_G_S20'].contiguous().detach().cpu(),self.batch_size)
            self.count_mean_G_S0D.update(records['loss_G_S0D'].contiguous().detach().cpu(),self.batch_size)
            self.count_mean_G_SD0.update(records['loss_G_SD0'].contiguous().detach().cpu(),self.batch_size)
            
            self.count_mean_D_S0.update(records["loss_D_S0"])
            self.count_mean_D_S1.update(records["loss_D_S1"])
            self.count_mean_D_S2.update(records["loss_D_S2"])
            self.count_mean_D_SD.update(records["loss_D_SD"])
            
            out["count_mean_G"]=self.count_mean_G.avg
            out["count_mean_G_S01"]=self.count_mean_G_S01.avg
            out["count_mean_G_S10"]=self.count_mean_G_S10.avg
            out["count_mean_G_S02"]=self.count_mean_G_S02.avg
            out["count_mean_G_S20"]=self.count_mean_G_S20.avg
            out["count_mean_G_S0D"]=self.count_mean_G_S0D.avg
            out["count_mean_G_SD0"]=self.count_mean_G_SD0.avg
            
            out["count_mean_D_S0"]=self.count_mean_D_S0.avg
            out["count_mean_D_S1"]=self.count_mean_D_S1.avg
            out["count_mean_D_S2"]=self.count_mean_D_S2.avg
            out["count_mean_D_SD"]=self.count_mean_D_SD.avg
            
        if self.controller == 'test' or not self.mode:
            if records['dualGAN']:
                loss_SB_PSNR = PSNR(records['fake_SAB'],records['real_SB'])
                loss_SAB_PSNR = PSNR(records['real_SA'],records['real_SB'])
                loss_SB_SSIM = SSIM(records['fake_SAB'],records['real_SB'])
                loss_SAB_SSIM = SSIM(records['real_SA'],records['real_SB'])
                
                self.count_mean_SB_PSNR.update(loss_SB_PSNR)
                self.count_mean_SAB_PSNR.update(loss_SAB_PSNR)
                self.count_mean_SB_SSIM.update(loss_SB_SSIM, self.batch_size, paths = records["image_paths"])
                self.count_mean_SAB_SSIM.update(loss_SAB_SSIM)
                    
                out["count_mean_SB_PSNR"] = self.count_mean_SB_PSNR.avg
                out["count_mean_SAB_PSNR"] = self.count_mean_SAB_PSNR.avg
                out["count_mean_SB_SSIM"] = self.count_mean_SB_SSIM.avg
                out["count_mean_SAB_SSIM"] = self.count_mean_SAB_SSIM.avg
                
                if self.use_inception:
                    
                    self.features['real_SA'].append(records['real_SA_Fea'].squeeze(0).squeeze(0).flatten().contiguous().cpu().detach().numpy())
                    self.features['real_SB'].append(records['real_SB_Fea'].squeeze(0).squeeze(0).flatten().contiguous().cpu().detach().numpy())
                    self.features['fake_SB'].append(records['fake_SB_Fea'].squeeze(0).squeeze(0).flatten().contiguous().cpu().detach().numpy())
                    
                    loss_SB_FID = self.loss(records['real_SB_Fea'],records['fake_SB_Fea'])
                    loss_SAB_FID = self.loss(records['real_SA_Fea'],records['real_SB_Fea'])
                    loss_DIS_FID = loss_SB_FID/(loss_SAB_FID+self.c)
                    self.count_mean_SB_FID.update(loss_SB_FID, self.batch_size, paths = records["image_paths"])
                    self.count_mean_SAB_FID.update(loss_SAB_FID)
                    self.count_mean_disFID.update(loss_DIS_FID, self.batch_size, paths = records["image_paths"])
                
                    out["count_mean_SB_FID"] = self.count_mean_SB_FID.avg
                    out["count_mean_SAB_FID"] = self.count_mean_SAB_FID.avg
                    out["count_mean_SB_DFID"] = self.count_mean_disFID.avg
                   
        # if self.frequency:
        loss_DF_SB = records["loss_Frequency_fake_SB"].contiguous().detach().cpu() / (records["loss_Frequency_real_SAB"].contiguous().detach().cpu() + self.c)
        self.count_DFrequency_SB.update(loss_DF_SB,self.batch_size, paths = records["image_paths"])
        self.count_Frequency_SB.update(records["loss_Frequency_fake_SB"].contiguous().detach().cpu(),self.batch_size, paths = records["image_paths"])
        out["count_Frequency_SB"] = change_fre(self.count_Frequency_SB.avg)

        return out
    
    def statistic(self):
        res = {}
        if self.mode:
            self.count_G.reset()
        else:
            # self.count_mean_S012D_SSIM.statistic()
            self.count_mean_SB_SSIM.statistic()

            if self.use_inception:
                # self.features = np.array(self.features)
                self.count_mean_SB_FID.statistic()
                self.count_mean_disFID.statistic()
               
                # var = np.var(self.features,axis=0)
                # x,y = self.features[:,np.argsort(var)[-1]],self.features[:,np.argsort(var)[-2]]
                
                
            if self.controller == 'test':
                # self.loss_mean_Frequency_SB = change_fre(self.count_Frequency_SB.avg)
                # if self.frequency:
                self.count_Frequency_SB.statistic()
                self.count_DFrequency_SB.statistic()
                self.count_DFrequency_SB.saveTop10p(os.path.join(os.path.abspath('.'),'results',self.name+'_v'+self.version,self.phase+'_'+self.epoch),'Top100p_DFre.txt')
                self.count_Frequency_SB.saveTop10p(os.path.join(os.path.abspath('.'),'results',self.name+'_v'+self.version,self.phase+'_'+self.epoch),'Top100p_Fre.txt')
                
                self.count_mean_SB_SSIM.saveTop10p(os.path.join(os.path.abspath('.'),'results',self.name+'_v'+self.version,self.phase+'_'+self.epoch),'Top100p_SSIM.txt')

                if self.use_inception:
                    self.count_mean_SB_FID.saveTop10p(os.path.join(os.path.abspath('.'),'results',self.name+'_v'+self.version,self.phase+'_'+self.epoch),'Top100p_FID.txt')
                    self.count_mean_disFID.saveTop10p(os.path.join(os.path.abspath('.'),'results',self.name+'_v'+self.version,self.phase+'_'+self.epoch),'Top100p_DFID.txt')
                    self.count_mean_SB_FID.reset()
                    self.count_mean_SAB_FID.reset()
                    # draw_Fea(np.array(self.features["real_SB"]),np.array(self.features["fake_SB"]),self.outputs,os.path.join(os.path.abspath('.'),'results',self.name+'_v'+self.version,self.phase+'_'+self.epoch,'last2feamap.png'))
                pass
            
            # self.TOP10p_mean_S012D_SSIM = self.count_mean_S012D_SSIM.Top10p
            # self.TOP50p_mean_S012D_SSIM = self.count_mean_S012D_SSIM.Top50p
            # self.TOP100p_mean_S012D_SSIM = self.count_mean_S012D_SSIM.avg
            res["TOP10p_mean_SB_SSIM"] = self.count_mean_SB_SSIM.Top10p
            res["TOP50p_mean_SB_SSIM"] = self.count_mean_SB_SSIM.Top50p
            
            self.count_mean_SB_PSNR.reset()
            self.count_mean_SAB_PSNR.reset()
            self.count_mean_SB_SSIM.reset()
            self.count_mean_SAB_SSIM.reset()
            # self.count_mean_S012D_SSIM.reset()

        # if self.frequency:
        self.count_Frequency_SB.reset()
            
                
        return res



